package com.iwindplus.base.mybatis.manager;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.text.CharSequenceUtil;
import com.baomidou.mybatisplus.autoconfigure.MybatisPlusProperties;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.core.toolkit.Constants;
import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.iwindplus.base.domain.DbSignBaseDO;
import com.iwindplus.base.domain.annotation.TableFieldSafe;
import com.iwindplus.base.domain.annotation.TableFieldSensitive;
import com.iwindplus.base.domain.enums.AlgorithmTypeEnum;
import com.iwindplus.base.mybatis.domain.property.MybatisProperty;
import com.iwindplus.base.mybatis.domain.property.MybatisProperty.FieldConfig.CryptoConfig;
import com.iwindplus.base.util.CryptoUtil;
import com.iwindplus.base.util.DbNameUtil;
import com.iwindplus.base.util.DbSignUtil;
import com.iwindplus.base.util.FilesUtil;
import com.iwindplus.base.util.JacksonUtil;
import com.iwindplus.base.util.SensitiveUtil;
import com.iwindplus.base.util.domain.dto.DbSignGenerateDTO;
import com.iwindplus.base.util.domain.dto.SensitiveDTO;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.experimental.SuperBuilder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.ReflectionUtils;

/**
 * MyBatis 字段级 加解密 & 脱敏 管理器（零反射运行期版） 支持 JDK 8/11/17+ 模块化，支持非 String 字段 JSON 序列化。 线程安全，可随配置热更新。
 *
 * @author zengdegui
 * @since 2025-04-22
 */
@Slf4j
@Getter
public class MybatisFieldCryptoManager {

    private static final Cache<Class<?>, EntityHolder> HOLDER_CACHE =
        Caffeine.newBuilder()
            .maximumSize(512)
            .expireAfterWrite(30, TimeUnit.MINUTES)
            .build();

    private static final Cache<String, SensitiveDTO> SENSITIVE_CACHE =
        Caffeine.newBuilder()
            .maximumSize(256)
            .build();

    private static final Cache<Object, Boolean> PROCESS_CACHE =
        Caffeine.newBuilder()
            .maximumSize(1024)
            .expireAfterWrite(10, TimeUnit.MINUTES)
            .build();

    private final Map<String, String> dbNameCache = new ConcurrentHashMap<>(16);
    private final Map<Class<?>, String> tableNameCache = new ConcurrentHashMap<>(16);

    private MybatisPlusProperties mybatisPlusProperties;
    private MybatisProperty mybatisProperty;
    private ThreadPoolTaskExecutor threadPoolTaskExecutor;

    public MybatisFieldCryptoManager(MybatisPlusProperties mybatisPlusProperties,
        MybatisProperty mybatisProperty, ThreadPoolTaskExecutor threadPoolTaskExecutor) {
        this.mybatisPlusProperties = mybatisPlusProperties;
        this.mybatisProperty = mybatisProperty;
        this.threadPoolTaskExecutor = threadPoolTaskExecutor;
        scanClassesField(mybatisPlusProperties);
        log.info("MybatisFieldManager 加载完成");
    }

    /**
     * 入库加密.
     *
     * @param source 数据源
     */
    public void encryptOnInput(Object source) {
        processHandler(source, true, false, false);
    }

    /**
     * 出库解密.
     *
     * @param source 数据源
     */
    public void decryptOnOutput(Object source) {
        processHandler(source, false, false, true);
    }

    /**
     * 入库脱敏.
     *
     * @param source 数据源
     */
    public void sensitiveOnInput(Object source) {
        processHandler(source, false, true, false);
    }

    /**
     * 出库脱敏.
     *
     * @param source 数据源
     */
    public void sensitiveOnOutput(Object source) {
        processHandler(source, false, true, false);
    }

    /**
     * 清空缓存.
     */
    public void clearCache() {
        HOLDER_CACHE.invalidateAll();
        SENSITIVE_CACHE.invalidateAll();
        PROCESS_CACHE.invalidateAll();

        log.info("MybatisFieldManager 缓存已清空");
    }

    /**
     * 设置签名.
     *
     * @param ms        ms
     * @param param     参数
     * @param secretKey 加签密钥
     * @param cmd       操作
     */
    public void addSign(MappedStatement ms, Object param, String secretKey, SqlCommandType cmd) {
        final List<Object> list = this.parseObj(param);
        if (CollUtil.isEmpty(list)) {
            return;
        }
        list.parallelStream().filter(Objects::nonNull).forEach(obj -> this.doSign(ms, param, secretKey, cmd));
    }

    /**
     * 解析对象.
     *
     * @param source 数据源
     * @return List
     */
    public List<Object> parseObj(Object source) {
        if (source == null) {
            return Collections.emptyList();
        }

        List<Object> result = new ArrayList<>(10);
        if (source instanceof Map<?, ?> map && MapUtil.isNotEmpty(map)) {
            map.values().forEach(obj -> result.addAll(parseObj(obj)));
        } else if (source instanceof Collection<?> coll && CollUtil.isNotEmpty(coll)) {
            coll.stream().filter(Objects::nonNull).forEach(obj -> result.addAll(parseObj(obj)));
        } else if (source instanceof Object[] array) {
            Arrays.stream(array).filter(Objects::nonNull).forEach(obj -> result.addAll(parseObj(obj)));
        } else {
            result.add(source);
        }

        return result;
    }

    /**
     * 获取数据库名.
     *
     * @param ms ms
     * @return String
     */
    public String getDbName(MappedStatement ms) {
        String envId = ms.getConfiguration().getEnvironment().getId();
        return dbNameCache.computeIfAbsent(envId, k -> {
            try {
                final Connection connection = ms.getConfiguration().getEnvironment().getDataSource().getConnection();
                if (connection == null) {
                    return null;
                }
                return DbNameUtil.getRealDbName(connection);
            } catch (SQLException e) {
                log.error("Get db name error, envId={}", envId, e);
                return null;
            }
        });
    }

    /**
     * 获取表名.
     *
     * @param clazz 类
     * @return 表名
     */
    public String getTableName(Class<?> clazz) {
        return tableNameCache.computeIfAbsent(clazz, k -> {
            TableInfo info = TableInfoHelper.getTableInfo(k);
            if (info == null) {
                return null;
            }
            return info.getTableName()
                .replace("`", "")
                .replace("\"", "");
        });
    }

    /**
     * 获取参数对象（用于输入时用）.
     *
     * @param paramObj
     * @return Object
     */
    public Object getParam(Object paramObj) {
        // 如果是 Map（含 ParamMap）
        if (paramObj instanceof Map<?, ?> map) {
            return map.get(Constants.ENTITY);
        } else {
            // 原生实体
            return paramObj;
        }
    }

    /**
     * 处理助手.
     *
     * @param source         原对象
     * @param inputEncrypt   是否启用输入加密
     * @param inputSensitive 是否启用输入脱敏
     * @param outputDecrypt  是否启用输出解密
     */
    public void processHandler(Object source, boolean inputEncrypt, boolean inputSensitive, boolean outputDecrypt) {
        final List<Object> list = parseObj(source);
        if (CollUtil.isEmpty(list)) {
            return;
        }

        if (list.size() >= this.mybatisProperty.getField().getCrypto().getBatchThreshold()) {
            final CompletableFuture[] completableFutures = list.parallelStream()
                .filter(Objects::nonNull)
                .map(o -> CompletableFuture.runAsync(() ->
                    processSingle(o, inputEncrypt, inputSensitive, outputDecrypt), threadPoolTaskExecutor))
                .toArray(CompletableFuture[]::new);
            CompletableFuture.allOf(completableFutures).join();
            return;
        }

        list.parallelStream().filter(Objects::nonNull)
            .forEach(o -> processSingle(o, inputEncrypt, inputSensitive, outputDecrypt));
        return;
    }

    private void processSingle(Object o, boolean enc, boolean sen, boolean dec) {
        // 检查是否已经处理
        int oHash = System.identityHashCode(o);
        if (PROCESS_CACHE.getIfPresent(oHash) != null) {
            log.debug("{} 已处理", o.getClass().getSimpleName());
            return;
        }

        final Class<?> clazz = o.getClass();
        final EntityHolder h = HOLDER_CACHE.get(clazz, k -> buildHolder(k));
        if (Objects.isNull(h)) {
            return;
        }

        try {
            if (enc) {
                apply(o, h.encrypt, this::doEncrypt);
            }
            if (dec) {
                apply(o, h.decrypt, this::doDecrypt);
            }
            if (sen) {
                apply(o, h.sensitive, this::doSensitive);
            }
        } catch (Throwable ex) {
            log.warn("process entity error, class={}", clazz.getSimpleName(), ex);
        } finally {
            PROCESS_CACHE.put(oHash, true);
        }
    }

    private void apply(Object o, FieldHandle[] handles, ThrowingOperator operator) throws Throwable {
        for (FieldHandle h : handles) {
            Object oldVal = h.getter.invoke(o);
            if (oldVal == null) {
                continue;
            }
            Object newVal = operator.apply(oldVal, h.meta);
            if (newVal != null && !newVal.equals(oldVal)) {
                h.setter.invoke(o, newVal);
            }
        }
    }

    private Object doEncrypt(Object val, FieldMetaDTO m) {
        AlgorithmTypeEnum alg = m.getAlgorithm();
        if (val instanceof String str) {
            return CryptoUtil.encrypt(str, cryptoConfig(alg));
        }
        return CryptoUtil.encrypt(JacksonUtil.toJsonStr(val), cryptoConfig(alg));
    }

    private Object doDecrypt(Object val, FieldMetaDTO m) {
        if (!(val instanceof String str)) {
            return val;
        }
        String dec = CryptoUtil.decrypt(str, cryptoConfig(m.getAlgorithm()));
        Class<?> type = m.getField().getType();
        return type == String.class ? dec : JacksonUtil.parseObject(dec, type);
    }

    private Object doSensitive(Object val, FieldMetaDTO m) {
        SensitiveDTO dto = buildSensitiveDTO(m.getSensitiveAnn());
        if (null == dto) {
            return val;
        }

        if (val instanceof String str) {
            return SensitiveUtil.desensitized(str, dto);
        }
        return SensitiveUtil.desensitized(JacksonUtil.toJsonStr(val), dto);
    }

    private void scanClassesField(MybatisPlusProperties props) {
        String pkg = props.getTypeAliasesPackage();
        Class<?> superType = props.getTypeAliasesSuperType();
        Set<Class<?>> classes = FilesUtil.scanClasses(pkg, superType);
        if (CollUtil.isEmpty(classes)) {
            return;
        }

        classes.parallelStream().forEach(clz -> HOLDER_CACHE.get(clz, k -> buildHolder(k)));
    }

    private EntityHolder buildHolder(Class<?> clz) {
        List<Field> all = FieldUtils.getAllFieldsList(clz);
        List<FieldHandle> encrypt = new ArrayList<>();
        List<FieldHandle> decrypt = new ArrayList<>();
        List<FieldHandle> sensitive = new ArrayList<>();

        for (Field f : all) {
            FieldHandle h = buildFieldHandle(f, clz);
            if (h == null) {
                continue;
            }
            if (h.meta.needEncrypt()) {
                encrypt.add(h);
            }
            if (h.meta.needDecrypt()) {
                decrypt.add(h);
            }
            if (h.meta.needDesensitize()) {
                sensitive.add(h);
            }
        }
        if (encrypt.isEmpty() && decrypt.isEmpty() && sensitive.isEmpty()) {
            return null;
        }
        return new EntityHolder(
            encrypt.toArray(FieldHandle[]::new),
            decrypt.toArray(FieldHandle[]::new),
            sensitive.toArray(FieldHandle[]::new)
        );
    }

    private FieldHandle buildFieldHandle(Field f, Class<?> clazz) {
        FieldMetaDTO meta = buildMeta(f);
        if (meta == null) {
            return null;
        }

        MethodHandles.Lookup lookup;
        MethodHandle getter, setter;
        try {
            lookup = MethodHandles.privateLookupIn(clazz, MethodHandles.lookup());
            getter = lookup.unreflectGetter(f);
            setter = lookup.unreflectSetter(f);
        } catch (Throwable t) {
            ReflectionUtils.makeAccessible(f);
            try {
                getter = MethodHandles.lookup().unreflectGetter(f);
                setter = MethodHandles.lookup().unreflectSetter(f);
            } catch (Throwable t1) {
                log.warn("Skip field {}#{} : {}", clazz.getSimpleName(), f.getName(), t.toString());
                return null;
            }
        }

        return new FieldHandle(meta, getter, setter, f);
    }

    private FieldMetaDTO buildMeta(Field f) {
        TableFieldSafe safe = f.getAnnotation(TableFieldSafe.class);
        TableFieldSensitive sen = f.getAnnotation(TableFieldSensitive.class);
        // 卫语句，避免写反
        if (safe == null && sen == null) {
            return null;
        }
        AlgorithmTypeEnum alg = safe == null ? null : safe.algorithm();
        return new FieldMetaDTO(f, alg, safe, sen);
    }

    private CryptoConfig cryptoConfig(AlgorithmTypeEnum alg) {
        CryptoConfig origin = mybatisProperty.getField().getCrypto();
        CryptoConfig clone = BeanUtil.copyProperties(origin, CryptoConfig.class);
        clone.setAlgorithm(alg);
        return clone;
    }

    private SensitiveDTO buildSensitiveDTO(TableFieldSensitive ann) {
        if (ann == null) {
            return null;
        }
        String key = ann.type() + "_" + ann.startInclude() + "_" + ann.endReserve();
        return SENSITIVE_CACHE.get(key, k ->
            SensitiveDTO.builder()
                .type(ann.type())
                .startInclude(ann.startInclude())
                .endReserve(ann.endReserve())
                .build());
    }

    private void doSign(MappedStatement ms, Object param, String secretKey, SqlCommandType cmd) {
        if (CharSequenceUtil.isBlank(secretKey) || !(param instanceof DbSignBaseDO entity)) {
            return;
        }

        // 获取db
        String dbName = getDbName(ms);
        if (CharSequenceUtil.isBlank(dbName)) {
            return;
        }

        // 获取table
        String tableName = getTableName(entity.getClass());
        if (CharSequenceUtil.isBlank(tableName)) {
            log.warn("Signature skipped: tableName is empty, class={}", entity.getClass());
            return;
        }

        String action = cmd.name();

        long salt = IdWorker.getId();
        entity.setSalt(salt);

        DbSignGenerateDTO dto = DbSignGenerateDTO.builder()
            .secretKey(secretKey)
            .salt(salt)
            .dbName(dbName)
            .tableName(tableName)
            .action(action)
            .build();
        String sign = DbSignUtil.generateSign(dto);
        entity.setSign(sign);
    }

    @FunctionalInterface
    private interface ThrowingOperator {

        Object apply(Object val, FieldMetaDTO meta) throws Throwable;
    }

    @Data
    @SuperBuilder
    @AllArgsConstructor
    private static final class FieldHandle {

        private final FieldMetaDTO meta;
        private final MethodHandle getter;
        private final MethodHandle setter;
        private final Field field;
    }

    @Data
    @SuperBuilder
    @AllArgsConstructor
    private static final class EntityHolder {

        private final FieldHandle[] encrypt;
        private final FieldHandle[] decrypt;
        private final FieldHandle[] sensitive;
    }

    @Data
    @SuperBuilder
    @AllArgsConstructor
    private static class FieldMetaDTO {

        private final Field field;
        private final AlgorithmTypeEnum algorithm;
        private final TableFieldSafe safeAnn;
        private final TableFieldSensitive sensitiveAnn;

        /**
         * 是否需要加密
         *
         * @return
         */
        boolean needEncrypt() {
            return safeAnn != null && safeAnn.inputEncrypt();
        }

        /**
         * 是否需要解密
         *
         * @return
         */
        boolean needDecrypt() {
            return safeAnn != null && safeAnn.outputDecrypt();
        }

        /**
         * 是否需要脱敏
         *
         * @return
         */
        boolean needDesensitize() {
            return sensitiveAnn != null && sensitiveAnn.enabled();
        }
    }

}